import argparse  # 导入 argparse 模块，用于解析命令行参数
import torch  # 导入 PyTorch 库，用于深度学习模型的构建和训练
from torch import optim  # 导入优化器模块，如 AdamW 等
from pathlib import Path  # 导入 Path 模块，用于处理文件路径
from tqdm import tqdm  # 导入进度条工具，用于显示训练进度
from crnn import WakeWordCRNN  # 导入自定义的 CRNN 模型（WakeWordCRNN）
from dataset import WakeWordDataset  # 导入自定义的数据集类（WakeWordDataset）
from torch import nn  # 导入神经网络模块，包含损失函数等


def train(args):  # 定义训练函数，接收命令行参数 args
    # 初始化设备
    device = torch.device(args.device)  # 根据传入的参数选择设备（CPU 或 GPU）
    print(device)  # 打印当前使用的设备

    # 加载数据集
    train_dataset = WakeWordDataset(args.train_dir, augment=True)  # 加载训练数据集，并启用数据增强
    valid_dataset = WakeWordDataset(args.valid_dir, augment=False)  # 加载验证数据集，不启用数据增强

    train_loader = train_dataset.get_loader(args.batch_size, num_workers=4)  # 创建训练数据加载器
    valid_loader = valid_dataset.get_loader(args.batch_size, shuffle=False)  # 创建验证数据加载器，关闭打乱顺序

    # 初始化模型
    model = WakeWordCRNN(input_dim=39).to(device)  # 初始化 CRNN 模型，并将其移动到指定设备
    optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-5)  # 使用 AdamW 优化器，设置学习率和权重衰减
    scheduler = optim.lr_scheduler.OneCycleLR(  # 使用 OneCycleLR 学习率调度器
        optimizer,
        max_lr=1e-3,  # 最大学习率
        steps_per_epoch=len(train_loader),  # 每个 epoch 的步数
        epochs=args.epochs,  # 总 epoch 数
        pct_start=0.3  # 学习率上升阶段占总周期的比例
    )
    criterion = nn.CrossEntropyLoss(weight=torch.tensor([1.0, args.pos_weight]).to(device))  # 初始化交叉熵损失函数，设置正样本权重

    # 在训练循环开始前添加
    print("正在验证数据集...")  # 打印提示信息，检查数据集是否正常
    for f, l, p in tqdm(train_loader):  # 遍历训练数据加载器
        if f.shape[-1] != 100:  # 检查特征长度是否符合预期
            print(f"发现异常特征：{p}，形状：{f.shape}")  # 打印异常特征的路径和形状
            raise ValueError("特征长度不一致，请检查音频预处理逻辑")  # 抛出异常，提醒检查音频预处理逻辑

    # 训练循环
    best_acc = 0.0  # 初始化最佳验证准确率
    for epoch in range(args.epochs):  # 遍历每个 epoch
        model.train()  # 设置模型为训练模式
        train_loss = 0.0  # 初始化训练损失
        correct = 0  # 初始化训练正确样本数
        total = 0  # 初始化训练总样本数

        # 训练阶段
        for features, labels, _ in tqdm(train_loader, desc=f"Epoch {epoch + 1}"):  # 遍历训练数据加载器
            features = features.to(device)  # 将特征移动到指定设备
            labels = labels.to(device)  # 将标签移动到指定设备

            optimizer.zero_grad()  # 清空优化器梯度
            outputs = model(features)  # 前向传播，计算模型输出
            loss = criterion(outputs, labels)  # 计算损失值
            loss.backward()  # 反向传播，计算梯度
            torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)  # 梯度裁剪，防止梯度爆炸
            optimizer.step()  # 更新模型参数

            train_loss += loss.item()  # 累加训练损失
            _, predicted = outputs.max(1)  # 获取预测结果
            correct += predicted.eq(labels).sum().item()  # 统计正确样本数
            total += labels.size(0)  # 累加总样本数

        # 验证阶段
        model.eval()  # 设置模型为评估模式
        val_loss = 0.0  # 初始化验证损失
        val_correct = 0  # 初始化验证正确样本数
        val_total = 0  # 初始化验证总样本数
        with torch.no_grad():  # 禁用梯度计算，节省内存
            for features, labels, _ in valid_loader:  # 遍历验证数据加载器
                features = features.to(device)  # 将特征移动到指定设备
                labels = labels.to(device)  # 将标签移动到指定设备

                outputs = model(features)  # 前向传播，计算模型输出
                loss = criterion(outputs, labels)  # 计算损失值

                val_loss += loss.item()  # 累加验证损失
                _, predicted = outputs.max(1)  # 获取预测结果
                val_correct += predicted.eq(labels).sum().item()  # 统计正确样本数
                val_total += labels.size(0)  # 累加总样本数

        # 计算指标
        train_acc = 100. * correct / total  # 计算训练准确率
        val_acc = 100. * val_correct / val_total  # 计算验证准确率
        scheduler.step(val_acc)  # 更新学习率调度器，基于验证准确率调整学习率

        # 保存最佳模型
        if val_acc > best_acc:  # 如果当前验证准确率高于历史最佳
            best_acc = val_acc  # 更新最佳验证准确率
            torch.save(model.state_dict(), Path(args.save_dir) / "best_model.pth")  # 保存模型权重

        # 打印日志
        print(f"Epoch {epoch + 1}/{args.epochs}")  # 打印当前 epoch 和总 epoch
        print(f"Train Loss: {train_loss / len(train_loader):.4f} | Acc: {train_acc:.2f}%")  # 打印训练损失和准确率
        print(f"Val Loss: {val_loss / len(valid_loader):.4f} | Acc: {val_acc:.2f}%")  # 打印验证损失和准确率
        print(f"LR: {optimizer.param_groups[0]['lr']:.2e}\n")  # 打印当前学习率


if __name__ == "__main__":  # 主程序入口
    parser = argparse.ArgumentParser()  # 创建命令行参数解析器
    parser.add_argument("--train_dir", default="dataset/split_dataset/train")  # 添加训练数据目录参数
    parser.add_argument("--valid_dir", default="dataset/split_dataset/val")  # 添加验证数据目录参数
    parser.add_argument("--save_dir", default="./checkpoints")  # 添加模型保存目录参数
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")  # 添加设备参数
    parser.add_argument("--batch_size", type=int, default=40)  # 添加批量大小参数
    parser.add_argument("--epochs", type=int, default=5000)  # 添加训练轮数参数
    parser.add_argument("--lr", type=float, default=1e-4)  # 添加学习率参数
    parser.add_argument("--pos_weight", type=float, default=3.0)  # 添加正样本权重参数
    args = parser.parse_args()  # 解析命令行参数

    # 创建保存目录
    Path(args.save_dir).mkdir(parents=True, exist_ok=True)  # 确保模型保存目录存在

    train(args)  # 调用训练函数，传入解析后的参数
    #这段代码实现了一个完整的深度学习训练流程，包括数据加载、模型初始化、训练、验证和模型保存。它适用于唤醒词检测任务，使用了 CRNN 模型和交叉熵损失函数，结合了学习率调度器以提高训练效果